package com.yichen.codesandbox.service.impl;

import cn.hutool.core.date.StopWatch;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.resource.ResourceUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.dfa.FoundWord;
import cn.hutool.dfa.WordTree;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.command.ExecStartResultCallback;
import com.yichen.codesandbox.model.dto.ExecutionCodeRequest;
import com.yichen.codesandbox.model.dto.ExecutionCodeResponse;
import com.yichen.codesandbox.model.dto.ExecutionMessage;
import com.yichen.codesandbox.model.dto.JudgeInfo;
import com.yichen.codesandbox.service.CodeSandBox;
import com.yichen.codesandbox.utils.ProcessUtils;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * @Author: yichenwu
 * @Date:2024/6/28
 */

/**
 * Java Docker代码沙箱
 */
public class JavaDockerCodeSandBox implements CodeSandBox {
    private final static String GLOBAL_CODE_DIR_NAME = "tempCode";
    private final static String GLOBAL_JAVA_CLASS_NAME = "Main.java";
    private final static Boolean FIRST_INIT = true;
    private final static Long TIME_OUT = 5 * 1000L;


    public static void main(String[] args) {
        // 测试编译
        JavaDockerCodeSandBox javaNativeCodeSandBox = new JavaDockerCodeSandBox();
        ExecutionCodeRequest executionCodeRequest = new ExecutionCodeRequest();
        String code = ResourceUtil.readStr("code/Main.java", StandardCharsets.UTF_8);
        executionCodeRequest.setCode(code);
        executionCodeRequest.setInput(Arrays.asList("1 2", "3 4"));
        javaNativeCodeSandBox.executeCode(executionCodeRequest);
    }

    @Override
    public List<ExecutionCodeResponse> executeCode(ExecutionCodeRequest executionCodeRequest) {
        try {
            // 1. 把用户的代码保存为文件
            String userDir = System.getProperty("user.dir");
            String globalCodePathName = userDir + File.separator + GLOBAL_CODE_DIR_NAME;
            // 如果不存在该全局目录则创建
            if (!FileUtil.exist(globalCodePathName)) {
                FileUtil.mkdir(globalCodePathName);
            }
            // 将用户代码隔离存放
            String userCodeParentPath = globalCodePathName + File.separator + UUID.randomUUID();
            String userCodePath = userCodeParentPath + File.separator + GLOBAL_JAVA_CLASS_NAME;
            String code = executionCodeRequest.getCode();
            File userCodefile = FileUtil.writeString(code, userCodePath, StandardCharsets.UTF_8);

            // 2. 编译代码，得到 class 文件
            ExecutionMessage compileMessage = ProcessUtils.processCompile(userCodefile);
            System.out.println(compileMessage);

            // 3.  创建docker容器，上传编译文件,控制内存等安全管理
            DockerClient dockerClient = DockerClientBuilder.getInstance().build();
            // 3.1 拉取镜像：(只初始化一次)
            String image = "openjdk:8-alpine";
            if (FIRST_INIT) {
                PullImageCmd pullImageCmd = dockerClient.pullImageCmd(image);
                PullImageResultCallback pullImageResultCallback = new PullImageResultCallback() {
                    @Override
                    public void onNext(PullResponseItem item) {
                        System.out.println("开始下载镜像jdk8..." + item.getStatus());
                        super.onNext(item);
                    }
                };
                pullImageCmd.exec(pullImageResultCallback)
                        .awaitCompletion();
                System.out.println("download succeed!");
            }

            // 3.2 创建容器：
            CreateContainerCmd containerCmd = dockerClient.createContainerCmd(image);
            // 设置创建容器时的一些配置
            HostConfig hostConfig = new HostConfig();
            // 指定文件路径（Volumn） 映射，作用是把本地的文件同步到容器中，可以让容器访问
            hostConfig.setBinds(new Bind(userCodeParentPath, new Volume("/app")));
            // 设置内存 -》100M
            hostConfig.withMemory(100 * 1000 * 1000L);
            // cpu
            hostConfig.withCpuCount(1L);
            // Linux 自带的一些安全管理措施，比如 seccomp（Secure Computing Mode）是一个用于 Linux 内核的安全功能，
            // 它允许你限制进程可以执行的系统调用，从而减少潜在的攻击面和提高容器的安全性。
            String profileConfig = ResourceUtil.readUtf8Str("profile.json");
            hostConfig.withSecurityOpts(Arrays.asList("seccomp=" + profileConfig));

            CreateContainerResponse containerResponse = containerCmd
                    // 容器连接本地终端，在终端控制台可以看到输出，也可以在终端输入
                    .withAttachStdin(true)
                    .withAttachStdout(true)
                    .withAttachStderr(true)
                    // 创建一个交互终端
                    .withTty(true)
                    // 容器挂载目录
                    .withHostConfig(hostConfig)
                    // 限制网络资源
                    .withNetworkDisabled(true)
                    // 权限管理: 限制用户不能向 root 根目录写文件
                    .withReadonlyRootfs(true)
                    .exec();
            String containerId = containerResponse.getId();
            System.out.println("创建容器:" + containerResponse);
            System.out.println("容器id:" + containerId);

            // 3.3 查看所有容器状态： todo:容器状态
//            List<String> containers =  new ArrayList();
//            containers.add(containerId);
//            ListContainersCmd listContainersCmd = dockerClient.listContainersCmd().withIdFilter(containers);
//            List<Container> exec = listContainersCmd.exec();
//            for (Container container : exec) {
//                if (!container.getState().equals("Created")){ // Running
//                    throw new RuntimeException("创建容器失败...");
//                }
//                System.out.println(container);
//            }

            // 3.4 启动容器
            dockerClient.startContainerCmd(containerId).exec();

            List<ExecutionCodeResponse> executionCodeResponseList = new ArrayList();
            // 4. 容器执行代码
            for (String inputStr : executionCodeRequest.getInput()) {
                List<String> outputList = new ArrayList();
                String[] inputArgs = inputStr.split(" ");
                // 示例执行命令 docker exec keen_blackwell java -cp /app Main 1 3  // keen_blackwell 容器名
                String[] cmdBasic = new String[]{"java", "-cp", "/app", "Main"};
                String[] cmd = ArrayUtil.append(cmdBasic, inputArgs);
                ExecCreateCmdResponse execCreateCmdResponse = dockerClient.execCreateCmd(containerId)
                        // 先创建命令
                        .withCmd(cmd)
                        .withAttachStdin(true)
                        .withAttachStderr(true)
                        .withAttachStdout(true)
                        .exec();

                ExecutionMessage executionMessage = new ExecutionMessage();
                final String[] executeMessage = {executionMessage.getExecuteMessage()};
                final String[] executeErrorMessage = {executionMessage.getExecuteErrorMessage()};
                final Integer[] exitValue = {executionMessage.getExitValue()};
                final Boolean[] isTimeOut = {true}; // 是否超时标志
                ExecStartResultCallback execStartResultCallback = new ExecStartResultCallback() {
                    @Override
                    public void onNext(Frame frame) {
                        StreamType streamType = frame.getStreamType();
                        exitValue[0] = streamType.equals(StreamType.STDERR) ? 1 : 0;
                        // 执行错误
                        if (streamType.equals(StreamType.STDERR)) {
                            executeErrorMessage[0] = new String(frame.getPayload());
                            System.out.println("输出错误结果: " + executeErrorMessage[0]);
                        } else {
                            executeMessage[0] = new String(frame.getPayload());
                            System.out.println("输出结果: " + executeMessage[0]);
                        }
                        super.onNext(frame);
                    }

                    @Override
                    public void onComplete() {
                        // 如果正常执行结束，表示未超时
                        isTimeOut[0] = false;
                        super.onComplete();
                    }
                };

                // 获取内存(一定时间片里的最大内存)
                final Long[] memory = {0L};
                StatsCmd statsCmd = dockerClient.statsCmd(containerId);
                ResultCallback<Statistics> statisticsResultCallback = statsCmd.exec(new ResultCallback<Statistics>() {
                    @Override
                    public void onNext(Statistics object) {
                        System.out.println("占用内存:" + object.getMemoryStats().getUsage());
                        memory[0] = Math.max(memory[0], object.getMemoryStats().getUsage());
                        System.out.println("memory[0]: " + memory[0]);
                    }

                    @Override
                    public void onStart(Closeable closeable) {
                    }

                    @Override
                    public void onError(Throwable throwable) {
                    }

                    @Override
                    public void onComplete() {
                    }

                    @Override
                    public void close() throws IOException {
                    }
                });
                statsCmd.exec(statisticsResultCallback);


                // 获取运行时间
                StopWatch stopWatch = new StopWatch();
                stopWatch.start();
                dockerClient.execStartCmd(execCreateCmdResponse.getId())
                        .exec(execStartResultCallback) // 执行完毕后再执行回调函数
                        // 超时控制
                        .awaitCompletion(TIME_OUT, TimeUnit.MICROSECONDS); // 异步执行
                stopWatch.stop();
                long lastTaskTimeMillis = stopWatch.getLastTaskTimeMillis();
                statsCmd.close();

                if (isTimeOut[0]) {
                    throw new RuntimeException("程序执行超时");
                }


                // 5. 收集整理输出结果
                ExecutionCodeResponse executionCodeResponse = new ExecutionCodeResponse();
                executionMessage.setExecuteMessage(executeMessage[0]);
                executionMessage.setExecuteErrorMessage(executeErrorMessage[0]);
                executionMessage.setExitValue(exitValue[0]);
                executionMessage.setMemory(memory[0]);
                executionCodeResponse.setMessage(executionMessage.toString());
                executionCodeResponse.setStatus(exitValue[0]);
                JudgeInfo judgeInfo = new JudgeInfo();
                if (executionMessage.getExecuteMessage() != null) {
                    outputList.add(executionMessage.getExecuteMessage());
                    judgeInfo.setMessage(executionMessage.getExecuteMessage());
                } else {
                    outputList.add(executionMessage.getExecuteErrorMessage());
                    judgeInfo.setMessage(executionMessage.getExecuteErrorMessage());
                }

                judgeInfo.setMemory(memory[0]);
                judgeInfo.setTime(lastTaskTimeMillis);
                executionCodeResponse.setJudgeInfo(judgeInfo);
                executionCodeResponse.setOutput(outputList);
                executionCodeResponseList.add(executionCodeResponse);
                System.out.println(executionCodeResponse);
            }

        } catch (Exception e) {
            e.getMessage();
            e.printStackTrace();
            getErrorExecutionCodeResponse(e);
        }
        return null;
    }


    // 6. 错误处理，提升程序健壮性
    private List<ExecutionCodeResponse> getErrorExecutionCodeResponse(Throwable e) {
        List<ExecutionCodeResponse> executionCodeResponseList = new ArrayList();
        ExecutionCodeResponse executionCodeResponse = new ExecutionCodeResponse();
        executionCodeResponse.setMessage(e.getMessage());
        executionCodeResponse.setStatus(2);
        executionCodeResponse.setJudgeInfo(new JudgeInfo());
        executionCodeResponse.setOutput(new ArrayList<>());
        executionCodeResponseList.add(executionCodeResponse);
        return executionCodeResponseList;
    }
}
